#!/usr/bin/python

__author__ = "Jimmy Saw"
__copyright__ = "Copyright 2011"
__credits__ = "Jimmy Saw"
__email__ = "jimmysaw@gmail.com"
__desc__ = "Partially curates draft annotation results."

"""
TO DO: 
Usage: python auto_anno.py annofile.txt

"""

import sys
import re
from Bio.Blast import NCBIXML

#Pattern matcher for inconsistent definitions
m1 = re.compile('hypothetical protein.*', re.IGNORECASE)
m2 = re.compile('predicted protein.*', re.IGNORECASE)
m3 = re.compile('protein of unknown function.*', re.IGNORECASE)
m4 = re.compile('PREDICTED.*')
m5 = re.compile('.*conserved hypothetical.*', re.IGNORECASE)
m6 = re.compile('putative orphan protein')
m7 = re.compile('(ribosomal protein S.*)', re.IGNORECASE)
m8 = re.compile('(ribosomal protein L.*)', re.IGNORECASE)
m9 = re.compile('50S ribosomal protein L31 type B')
m10 = re.compile('ribosomal protein L11 methyltransferase')

annofile = sys.argv[1]
af = open(annofile, "rU")
lines = af.readlines()
num = len(lines)

crispr_starts = [141632, 3742956, 3820468]
crispr_stops = [153410, 3753501, 3828723]

i = 0
length_check = {}
fs_check = {}
overlap_check = {}
new_desc = {}
crispr_check = {}
cds_list = []
cds_starts = []
cds_stops = []
rna_list = []
rna_starts = []
rna_stops = []

while i < num:
    curr_line = lines[i].split('\t')
    curr_id = curr_line[0]
    curr_locus_tag = curr_line[1]
    curr_feat_type = curr_line[2]
    curr_start = int(curr_line[3])
    curr_stop = int(curr_line[4])
    curr_qlen = (curr_stop - curr_start + 1 - 3) / 3
    curr_frame = curr_line[5]
    curr_blast_status = curr_line[6]
    curr_tophit_acc = curr_line[10]
    curr_hitlength = curr_line[14]
    curr_hit_def = curr_line[17]

    if i == 0:
        next_line = lines[i + 1].split('\t')
        next_id = next_line[0]
        next_start = int(next_line[3])
        next_stop = int(next_line[4])
        next_blast_status = next_line[6]
        next_tophit_acc = next_line[10]
        if curr_feat_type == "CDS":
            cds_list.append(curr_locus_tag)
            cds_starts.append(curr_start)
            cds_stops.append(curr_stop)
        #Work on the coding sequences to check for pseudogenes
            if curr_blast_status == "Yes":
                if m1.match(curr_hit_def) or m2.match(curr_hit_def) or m3.match(curr_hit_def) or m4.match(curr_hit_def):
                    new_desc[curr_locus_tag] = "hypothetical protein"
                elif m5.match(curr_hit_def):
                    new_desc[curr_locus_tag] = "conserved hypothetical protein"
                elif m6.match(curr_hit_def):
                    new_desc[curr_locus_tag] = "hypothetical protein"
                else:
                    #Checking for adjacent genes
                    if curr_tophit_acc == next_tophit_acc:
                        fs_check[curr_locus_tag] = "Maybe frameshifted"
                #Checking for truncated or fused genes
                if int(curr_qlen) < (0.8 * int(curr_hitlength)):
                    length_check[curr_locus_tag] = "Possibly truncated" + "\t" + str(curr_qlen) + "<" + str(curr_hitlength)
                elif int(curr_qlen) > (1.2 * int(curr_hitlength)):
                    length_check[curr_locus_tag] = "Possibly fused" + "\t" + str(curr_qlen) + ">" + str(curr_hitlength)
            else:
                if curr_stop - curr_start + 1 < 90:
                    length_check[curr_locus_tag] = "Length is < 90bp. Remove!"
                if curr_stop >= next_start + 40:
                    overlap_length = curr_stop - next_start
                    if next_blast_status == "Yes":
                        overlap_check[curr_locus_tag] = "Overlaps with next CDS with BLAST hit. \
                        Overlap length = " + str(overlap_length) + ". Remove!"
    elif i == num - 1:
        prev_line = lines[i - 1].split('\t')
        prev_id = prev_line[0]
        prev_start = int(prev_line[3])
        prev_stop = int(prev_line[4])
        prev_blast_status = prev_line[6]
        prev_tophit_acc = prev_line[10]
        if curr_feat_type == "CDS":
            cds_list.append(curr_locus_tag)
            cds_starts.append(curr_start)
            cds_stops.append(curr_stop)
        #Work on the coding sequences to check for pseudogenes
            if curr_blast_status == "Yes":
                if m1.match(curr_hit_def) or m2.match(curr_hit_def) or m3.match(curr_hit_def) or m4.match(curr_hit_def):
                    new_desc[curr_locus_tag] = "hypothetical protein"
                elif m5.match(curr_hit_def):
                    new_desc[curr_locus_tag] = "conserved hypothetical protein"
                elif m6.match(curr_hit_def):
                    new_desc[curr_locus_tag] = "hypothetical protein"
                else:
                    #Checking for adjacent genes
                    if curr_tophit_acc == prev_tophit_acc:
                        fs_check[curr_locus_tag] = "Maybe frameshifted"
                #Checking for truncated or fused genes
                if int(curr_qlen) < (0.8 * int(curr_hitlength)):
                    length_check[curr_locus_tag] = "Possibly truncated" + "\t" + str(curr_qlen) + "<" + str(curr_hitlength)
                elif int(curr_qlen) > (1.2 * int(curr_hitlength)):
                    length_check[curr_locus_tag] = "Possibly fused" + "\t" + str(curr_qlen) + ">" + str(curr_hitlength)
            else:
                if curr_stop - curr_start + 1 < 90:
                    length_check[curr_locus_tag] = "Length is < 90bp. Remove!"
                if curr_start <= prev_stop - 40:
                    overlap_length = prev_stop - curr_start
                    if prev_blast_status == "Yes":
                        overlap_check[curr_locus_tag] = "Overlaps with previous CDS with BLAST hit. \
                        Overlap length = " + str(overlap_length) + ". Remove!"
    elif i != 0 and i != num - 1:
        next_line = lines[i + 1].split('\t')
        next_id = next_line[0]
        next_start = int(next_line[3])
        next_stop = int(next_line[4])
        next_blast_status = next_line[6]
        next_tophit_acc = next_line[10]
        prev_line = lines[i - 1].split('\t')
        prev_id = prev_line[0]
        prev_start = int(prev_line[3])
        prev_stop = int(prev_line[4])
        prev_blast_status = prev_line[6]
        prev_tophit_acc = prev_line[10]

        if curr_feat_type == "CDS":
            cds_list.append(curr_locus_tag)
            cds_starts.append(curr_start)
            cds_stops.append(curr_stop)
        #Work on the coding sequences to check for pseudogenes
            if curr_blast_status == "Yes":
                if m1.match(curr_hit_def) or m2.match(curr_hit_def) or m3.match(curr_hit_def) or m4.match(curr_hit_def):
                    new_desc[curr_locus_tag] = "hypothetical protein"
                elif m5.match(curr_hit_def):
                    new_desc[curr_locus_tag] = "conserved hypothetical protein"
                elif m6.match(curr_hit_def):
                    new_desc[curr_locus_tag] = "hypothetical protein"
                else:
                    #Checking for adjacent genes
                    if curr_tophit_acc == next_tophit_acc or curr_tophit_acc == prev_tophit_acc:
                        fs_check[curr_locus_tag] = "Maybe frameshifted"
                #Checking for truncated or fused genes
                if int(curr_qlen) < (0.8 * int(curr_hitlength)):
                    length_check[curr_locus_tag] = "Possibly truncated" + "\t" + str(curr_qlen) + "<" + str(curr_hitlength)
                elif int(curr_qlen) > (1.2 * int(curr_hitlength)):
                    length_check[curr_locus_tag] = "Possibly fused" + "\t" + str(curr_qlen) + ">" + str(curr_hitlength)
                x = 0
                y = len(crispr_starts)
                while x < y:
                    if crispr_starts[x] < curr_start < crispr_stops[x] or crispr_starts[x] < curr_stop < crispr_stops[x]:
                        crispr_check[curr_locus_tag] = "Overlaps with CRISPR regions. Remove!"
                    x += 1
            else:
                if curr_stop - curr_start + 1 < 90:
                    length_check[curr_locus_tag] = "Length is < 90bp. Remove!"
                if curr_stop >= next_start + 40:
                    overlap_length = curr_stop - next_start
                    if next_blast_status == "Yes":
                        overlap_check[curr_locus_tag] = "Overlaps with next CDS with BLAST hit. \
                        Overlap length = " + str(overlap_length) + ". Remove!"
                elif curr_start <= prev_stop - 40:
                    overlap_length = prev_stop - curr_start
                    if prev_blast_status == "Yes":
                        overlap_check[curr_locus_tag] = "Overlaps with previous CDS with BLAST hit. \
                        Overlap length = " + str(overlap_length) + ". Remove!"
                x = 0
                y = len(crispr_starts)
                while x < y:
                    if crispr_starts[x] < curr_start < crispr_stops[x] or crispr_starts[x] < curr_stop < crispr_stops[x]:
                        crispr_check[curr_locus_tag] = "Overlaps with CRISPR regions. Remove!"
                    x += 1
        elif curr_feat_type == "tRNA":
            rna_list.append(curr_hit_def)
            rna_starts.append(curr_start)
            rna_stops.append(curr_stop)
        elif curr_feat_type == "rRNA":
            rna_list.append(curr_hit_def)
            rna_starts.append(curr_start)
            rna_stops.append(curr_stop)
    i += 1

rna_overlaps = {}
#Checking RNA coords
for i, x in enumerate(rna_list):
    for j, k in enumerate(cds_list):
        if cds_starts[j] < rna_starts[i] < cds_stops[j] or cds_starts[j] < rna_stops[i] < cds_stops[j]:
            rna_overlaps[k] = "Overlaps with RNA region. Remove!"
#Checking CDS coords
for i, x in enumerate(cds_list):
    for j, k in enumerate(rna_list):
        if rna_starts[j] < cds_starts[i] < rna_stops[j] or rna_starts[j] < cds_stops[i] < rna_stops[j]:
            rna_overlaps[x] = "Overlaps with RNA region. Remove!"

for line in lines:
    r = line.split('\t')
    pid = r[0]
    locus_tag = r[1]
    feat = r[2]
    start = r[3]
    stop = r[4]
    aa_len = (int(stop) - int(start) + 1 - 3) / 3
    frame = r[5]
    desc = r[17]
    cog_id = r[26]
    cog_cat = r[27]
    cog = cog_id + cog_cat
    gene_name = r[30]
    cluster_def = r[31]
    new_description = ""
    length_result = ""
    fs_result = ""
    overlap_result = ""
    has_cluster = ""
    crispr_result = ""
    pfam_hits = ""
    pattern = ""
    
    #Fixing annotation definitions
    if new_desc.has_key(locus_tag):
        new_description = new_desc[locus_tag]
    else:
        if cluster_def != "Not in cluster" and cluster_def != "hypothetical protein":
            has_cluster = "Cluster Definition"
            if m10.match(cluster_def):
                new_description = cluster_def
            elif m7.match(cluster_def):
                pattern = m7.match(cluster_def)
                new_description = "30S " + pattern.group(1)
            elif m8.match(cluster_def):
                pattern = m8.match(cluster_def)
                new_description = "50S " + pattern.group(1)
            else:
                new_description = cluster_def
        else:
            has_cluster = ""
            if m10.match(desc):
                new_description = desc
            elif m7.match(desc):
                pattern = m7.match(desc)
                new_description = "30S " + pattern.group(1)
            elif m8.match(desc):
                pattern = m8.match(desc)
                new_description = "50S " + pattern.group(1)
            else:
                new_description = desc
    if m9.match(desc):
        new_description = "hypothetical protein"

    if length_check.has_key(locus_tag):
        length_result = length_check[locus_tag]
        pfam_xmlfile = pid + ".rpsblast.Pfam.xml"
        pfam_handle = open(pfam_xmlfile)
        pfam = NCBIXML.parse(pfam_handle)
        pfam_rec = pfam.next()
        hsp_list = []
       
        for hsps in pfam_rec.descriptions:
            hsp_list.append(hsps.title)
        for alignment in pfam_rec.alignments:
            for hsp in alignment.hsps:
                if alignment.length * 0.9 < hsp.align_length < alignment.length * 1.1:
                    length_result = length_check[locus_tag] + "\t" + "Has Pfam hit"
                else:
                    length_result = length_check[locus_tag]
    else:
        length_result = ""
    if fs_check.has_key(locus_tag):
        fs_result = fs_check[locus_tag]
    else:
        fs_result = ""
    if overlap_check.has_key(locus_tag):
        overlap_result = overlap_check[locus_tag]
    else:
        overlap_result = ""
    if crispr_check.has_key(locus_tag):
        crispr_result = crispr_check[locus_tag]
    else:
        crispr_result = ""

    if feat == "CDS":
        #This will remove all the bad CDS (hypotheticals) annotated as overlaping with good regions
        if locus_tag not in overlap_check and locus_tag not in crispr_check and locus_tag not in rna_overlaps:
            print feat + "\t" + pid + "\t" + locus_tag + "\t" + start + "\t" + stop + "\t" + str(aa_len) + "\t" \
                + frame + "\t" + gene_name + "\t" + cog + "\t" + desc + "\t" + new_description + "\t" + has_cluster + \
                "\t" + length_result + "\t" + fs_result + "\t" + overlap_result + "\t" + crispr_result
    elif feat == "tRNA":
        print feat + "\t" + "None" + "\t" + locus_tag + "\t" + start + "\t" + stop + "\t" + "" + "\t" + \
            frame + "\t" + "" + "\t" + "" + "\t" + desc
    elif feat == "rRNA":
        print feat + "\t" + "None" + "\t" + locus_tag + "\t" + start + "\t" + stop + "\t" + "" + "\t" + \
            frame + "\t" + "" + "\t" + "" + "\t" + desc
        
af.close()
